import torch
from build_gym import BipedalWheeledRobotEnv 
from PPO_2 import Agent
import numpy as np
 

# Create the environment
env = BipedalWheeledRobotEnv(render_mode='human')  

# Create the agent
agent = Agent()

EP_MAX = 10000
HORIZON = 200

# Training loop
for episode in range(EP_MAX):
    state, _ = env.reset()
    episode_reward = 0

    for t in range(HORIZON):
        # Convert state to tensor and select action
        state_tensor = torch.tensor(state, dtype=torch.float32)
        action = agent.choose_action(state_tensor)

        # Execute action, get new state, reward, and done flag
        next_state, reward, done, truncated, info = env.step(np.array([action]))
        agent.push_data((state, action, reward, next_state, done))

        # Update current state
        state = next_state
        episode_reward += reward

        if done or truncated:
            break

    # Update the agent
    agent.updata()

    # Print reward for each episode
    print(f"Episode: {episode + 1}, Reward: {episode_reward}")

# Close the environment
env.close()
